from .cgc_exploit import CGCExploit
from .c_templates import c_template_type2


class CGCType2Exploit(CGCExploit):
    '''
    A CGC exploit object, offers more flexibility than an Exploit object for
    the sake of the game. This should represent a Type 2 POV:
    https://github.com/CyberGrandChallenge/cgc-release-documentation/blob/master/walk-throughs/understanding-cfe-povs.md
    '''

    def __init__(self, crash, bypasses_nx, bypasses_aslr):
        '''
        :param crash: a crash object which has been modified to exploit a
        vulnerability
        :param bypasses_nx: does the exploit bypass NX?
        :param bypasses_aslr: does the exploit bypass ASLR?
        '''
        super(CGCType2Exploit, self).__init__(crash, 2, bypasses_nx,
                                              bypasses_aslr)

        self._output_leak_idx = None

    def dump_c(self, filename=None):
        """
        Creates a simple C file to do the type1 exploit
        :param filename: dumps the code to this path if filename is not None
        :return: the c_code
        """
        encoded_payload = ""
        for c in self._raw_payload:
            encoded_payload += "\\x%02x" % c

        fmt_args = dict()
        fmt_args["payload_len"] = hex(self._payload_len)
        fmt_args["raw_payload"] = encoded_payload
        fmt_args["output_leak_idx"] = hex(self._output_leak_idx)
        fmt_args["solver_code"] = self._solver_code
        fmt_args["recv_buf_len"] = hex(self._recv_buf_len)
        # int stuff
        fmt_args["payload_int_start_locations"] = self._make_c_int_arr([x.start for x in self._sorted_stdin_int_infos])
        fmt_args["payload_int_bases"] = self._make_c_int_arr([x.base for x in self._sorted_stdin_int_infos])
        fmt_args["payload_int_expected_lens"] = self._make_c_int_arr([x.size for x in self._sorted_stdin_int_infos])
        fmt_args["recv_int_start_locations"] = self._make_c_int_arr([x.start for x in self._sorted_stdout_int_infos])
        fmt_args["recv_int_bases"] = self._make_c_int_arr([x.base for x in self._sorted_stdout_int_infos])
        fmt_args["recv_int_expected_lens"] = self._make_c_int_arr([x.size for x in self._sorted_stdout_int_infos])
        fmt_args["num_payload_ints"] = str(len(self._sorted_stdin_int_infos))
        fmt_args["num_recv_ints"] = str(len(self._sorted_stdout_int_infos))

        # TODO using .format is annoying because of all the curly braces
        # figure out how to do this better
        c_code = c_template_type2.c_template
        for k, v in fmt_args.items():
            c_code = c_code.replace("{%s}" % k, v)

        if filename is not None:
            with open(filename, 'w') as f:
                f.write(c_code)
        else:
            return c_code
